# -*- coding: utf-8 -*-
# @Author   : SongLiangCheng
# @QQ       : 2192383945
# @Time     : 2022/11/29 8:48
# @File     : http.py
# @Project  : scraping
# @Desc     :  通用爬取方法
# pip install coloredlogs openai redis
import json
import logging
import threading
import time
import traceback
from concurrent.futures._base import Future
from concurrent.futures.thread import ThreadPoolExecutor
from dataclasses import dataclass

import coloredlogs
import openai
import redis


@dataclass
class Configuration:
    # redis
    redis_conn: 'MyRedis' = None
    redis_host: str = 'localhost'
    redis_port: int = 6379
    redis_db: int = 0
    redis_password: str = 'mykernel.cn'

    # OpenAI
    api_key: str = ''
    # 4096 - prompt.length
    max_tokens: int = 3900

    # threads
    threads: int = 11
    formatter_string: str = '%(asctime)s - %(threadName)s : {%(name)s:%(lineno)d} : %(levelname)s : %(message)s'
    debug: bool = True

    def load_json(self):
        with open('config.json', 'r') as fb:
            config = json.load(fb)
        for field_name in self.__dataclass_fields__:
            if field_name in config:
                setattr(self,field_name,config[field_name])
    def show_config(self):
        print(json.dumps({k: getattr(self,k) for k,v in self.__dataclass_fields__.items()},indent=4))

@dataclass
class MyRedis():
    config: Configuration

    def __post_init__(self):
        self.origin_redis = redis.Redis(host=config.redis_host, port=config.redis_port,
                                        db=config.redis_db, password=config.redis_password)

    def lpush(self, value, name='msg'):
        return self.origin_redis.execute_command("LPUSH", name, value)

    def rpop(self, name='msg'):
        return self.origin_redis.execute_command("RPOP", name)


@dataclass
class OpenAI:
    config: Configuration

    def run(self, input):
        """
        model: https://beta.openai.com/docs/models/gpt-3
        :param input:
        :return:
        """
        openai.api_key = self.config.api_key
        res = openai.Completion.create(
            model="text-davinci-003",
            prompt=input,
            temperature=1,
            max_tokens=self.config.max_tokens,
        )
        return res.choices[0].text


@dataclass
class Daemon:
    config: Configuration

    def work(self, prompt: str):
        try:
            openapi = OpenAI(config=self.config)
            text = openapi.run(prompt)
            logger.info(f"\n{prompt}\n========================\n{text}\n\n")
        except:
            logger.error(f"发现异常 {str(traceback.format_exc())}")

    def _threadpool(self):
        pool = ThreadPoolExecutor(max_workers=self.config.threads)
        daemons: list[Future] = []

        while True:
            item = self.config.redis_conn.rpop()
            logger.debug(f"获取到消息{item}")
            if not item:
                time.sleep(2)
                logger.debug('没有消息休息2s')
                continue
            item = item.decode()
            logger.warning(f'拿到消息: --> {item}')

            daemons.append(pool.submit(self.work, item))  # 线程中启动多个线程

    def threadpool(self):
        threading.Thread(target=self._threadpool).start()

    def run(self):
        # 线程
        self.threadpool()


def log_config(config):
    formatter = logging.Formatter(config.formatter_string)

    # 定义root handler
    coloredlogs.install(fmt=config.formatter_string)  # 默认INFO

    # 定义logger
    logger = logging.getLogger(__file__)

    # 定义debug
    coloredlogs.install(fmt=config.formatter_string, level=logging.DEBUG if config.debug else logging.INFO)
    return logger


if __name__ == '__main__':
    config = Configuration()
    config.load_json() # 配置优先
    config.show_config() # 打印生效的配置
    config.redis_conn = MyRedis(config=config)

    logger = log_config(config)

    daemons: list[Daemon] = []

    # 1个daemon
    daemons.append(Daemon(config=config))

    try:
        logger.info('开始启动....')
        for daemon in daemons:
            daemon.run()
        while True:
            logger.debug(threading.enumerate())
            time.sleep(2)
    except:
        logger.info('退出')
        exit(1)
